Fix heap out-of-bounds write in MaxPoolGrad via unchecked indices#27932
Closed
vraspar wants to merge 2 commits intomicrosoft:mainfrom
Closed
Fix heap out-of-bounds write in MaxPoolGrad via unchecked indices#27932vraspar wants to merge 2 commits intomicrosoft:mainfrom
vraspar wants to merge 2 commits intomicrosoft:mainfrom
Conversation
Add bounds validation for index values in MaxPoolGrad::Compute to prevent heap out-of-bounds writes when the indices tensor contains values outside the valid range [0, output_size). Each index is now validated with ORT_ENFORCE before being used as an offset into the output buffer. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
5625230 to
2d90937
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens the training MaxPoolGrad CPU kernel against malicious or malformed indices tensors by adding explicit bounds validation before writing into the dX gradient buffer, and adds unit tests to cover negative and out-of-range indices.
Changes:
- Add per-element range checks for
indicesinMaxPoolGrad::Computeand simplify the write path (dX_data[index] += ...). - Add CPU unit tests that verify failure on negative and out-of-range
indicesvalues.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| orttraining/orttraining/training_ops/cpu/nn/pool_gradient_op.cc | Adds index bounds validation to prevent heap out-of-bounds writes in MaxPoolGrad. |
| orttraining/orttraining/test/training_ops/cpu/nn/pool_gradient_op_test.cc | Adds negative/out-of-range indices regression tests for MaxPoolGrad. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
orttraining/orttraining/training_ops/cpu/nn/pool_gradient_op.cc
Outdated
Show resolved
Hide resolved
Return INVALID_ARGUMENT status instead of aborting to avoid DoS when built with ORT_NO_EXCEPTIONS, where ORT_ENFORCE maps to abort(). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Author
|
Closing in favor of #27903 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds bounds validation for index values in
MaxPoolGrad::Computeto prevent heap out-of-bounds writes when theindicestensor contains values outside the valid range[0, output_size).Motivation
The
MaxPoolGradoperator uses values from theindicesinput tensor as direct offsets into the output gradient buffer (dX_data) without validating they are within bounds. A maliciously craftedindicestensor with negative or out-of-range values can write to arbitrary heap memory.Changes
orttraining/orttraining/training_ops/cpu/nn/pool_gradient_op.ccdX_shape.Size()intodX_sizeto avoid repeated callsindices_data[i]is in[0, dX_size)usingORT_ENFORCEThis follows the same pattern used in the recent RoiAlign OOB fix (#27543).